import time

import pypdfium2 # Needs to be at the top to avoid warnings
import os
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # For some reason, transformers decided to use .isin for a simple op, which is not supported on MPS

import argparse
from marker.convert import convert_single_pdf
from marker.logger import configure_logging
from marker.models import load_all_models

from marker.output import save_markdown


import os

from joblib import dump, load

Model_Variable_Filepath = '/Users/Shared/model_cache/marker-pdf2md/marker-model-big_variable.joblib'

def load_exists_model_list(model_big_filepath = Model_Variable_Filepath ):
    if  os.path.exists(model_big_filepath):
        print('加载已经存在的模型')
        # 加载变量
        loaded_variable = load(model_big_filepath)
        return loaded_variable
    # 加载模型
    model_lst = load_all_models()
    # 获取文件的目录路径
    directory = os.path.dirname(model_big_filepath)
    #
    # 如果目录不存在，创建目录
    if directory and not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Created directory: {directory}")
    #
    # 保存变量
    dump(model_lst, model_big_filepath)
    return model_lst


model_big_filepath = Model_Variable_Filepath
model_lst = load_exists_model_list(model_big_filepath )



configure_logging()




class ArgsInput:
    max_pages = 10
    batch_multiplier = 1
    start_page = 10
    ocr_all_pages = False
    langs = None
    filename= '/Users/emery/2-Pre/UCAS-ppt-slide/政治-中国特色理论实践/zhongte.pdf'
    output = "/Users/emery/Downloads/pdf2markdown-output"


args = ArgsInput()

langs = args.langs.split(",") if args.langs else None
fname = args.filename
start = time.time()

full_text, images, out_meta = convert_single_pdf(fname, model_lst,
    max_pages=args.max_pages,
    langs=langs,
    batch_multiplier=args.batch_multiplier,
    start_page=args.start_page,
    ocr_all_pages=args.ocr_all_pages)




fname = os.path.basename(fname)
fname_output =  f"{fname}-{args.start_page}"
subfolder_path = save_markdown(args.output,
    fname_output, full_text, images, out_meta)




print(f"Saved markdown to the {subfolder_path} folder")
print(f"Total time: {time.time() - start}")


#

import os
from marker.output import save_markdown
import random
import string

def generate_random_string(length=5):
    # 字符集包括大小写字母和数字
    characters = string.ascii_letters + string.digits
    return ''.join(random.choice(characters) for _ in range(length))


def save_markdown_to_dir(full_text, images, out_meta,
    filename, output_folder= None):
    output_dir = output_folder
    if  output_folder is None :
        user_home = os.path.expanduser("~")
        output_dir = os.path.join(user_home, "Downloads")
    #
    random_string = generate_random_string()
    fname_output =  f"{filename}-{random_string}"
    subfolder_path = save_markdown(output_dir,
        fname_output, full_text, images, out_meta)



random_string = generate_random_string()

# 获取当前用户的主目录
user_home = os.path.expanduser("~")
download_dir = os.path.join(user_home, "Downloads")
print(f"当前用户的主目录是: {user_home}")





print(random_string)



md_text, images, out_meta



print(f"当前用户的主目录是: {user_home}")

# 如果你想要获取特定的子目录，可以这样做：
documents_dir = os.path.join(user_home, "Documents")
desktop_dir = os.path.join(user_home, "Desktop")
download_dir = os.path.join(user_home, "Downloads")
print(f"文档目录: {documents_dir}")
print(f"桌面目录: {desktop_dir}")



#
